import unittest
from notebook.image_utils import *


class MyTestCase(unittest.TestCase):

    def setUp(self) -> None:
        data_root_orig = 'E:\\Datasets\\Image Classification\\tiny-imagenet-200'
        self.data_root = Path(data_root_orig)
        print(self.data_root)

    def test_retrieve_images(self):
        """
        检索图片
        """
        for item in self.data_root.iterdir():
            print(item)

        train_image_count = count_image(self.data_root, 'train/*/images/*')
        self.assertEqual(100000, train_image_count)

        val_image_count = count_image(self.data_root, 'test/images/*')
        self.assertEqual(10000, val_image_count)

        test_image_count = count_image(self.data_root, 'val/images/*')
        self.assertEqual(10000, test_image_count)

        val_image_paths = get_all_image_paths(self.data_root, 'val/images/*')
        print(val_image_paths[:10])
